#pragma once

struct CExportedFunction
{
	explicit CExportedFunction(unsigned entryPoint, unsigned ordinal, LPCSTR name, LPCSTR forwarder)
		: m_EntryPoint(entryPoint)
		, m_Ordinal(ordinal)
	{
		if (name)
		{
			m_Name = name;
		}

		if (forwarder != NULL)
		{
			m_Forwarder = forwarder;
		}
	}

	unsigned       Ordinal()    const { return m_Ordinal;    }
	unsigned       EntryPoint() const { return m_EntryPoint; }
	const string&  Name()       const { return m_Name;       }
	const string&  Forwarder()  const { return m_Forwarder;  }

private:
	unsigned m_EntryPoint;
	unsigned m_Ordinal;

	string  m_Name;
	string  m_Forwarder;
};

class CPEImage
{
public:
	typedef vector<CExportedFunction> export_vector;

	explicit CPEImage(const CAtlFile& file);
	explicit CPEImage(const tstring& path);

	bool IsManagedAssembly() const { return m_Managed; }
	bool IsWin32PE()         const { return m_Win32 && !m_Managed; }
	bool IsComServer()       const { return m_COM; }
	const export_vector& ExportedFunctions() const
	{
		return m_Exports;
	}

private:

	void Initialize();

	template<class T>
	T MakePtr(PBYTE base, ULONGLONG addend) const
	{
		return reinterpret_cast<T>(base + addend);
	}

	template<class THeader>
	const IMAGE_DATA_DIRECTORY* GetDirectory(unsigned index, const THeader* pNTHeader) const
	{
		if (pNTHeader->OptionalHeader.NumberOfRvaAndSizes <= index)
		{
			return NULL;
		}

		return &(pNTHeader->OptionalHeader.DataDirectory[index]);
	}

	
	template <class TPtr, class THeader>
	TPtr* GetPtrFromRVA(DWORD rva, const THeader* pNTHeader) const
	{
		return reinterpret_cast<TPtr*>(GetPtrFromRVA(rva, pNTHeader));
	}


	template<class THeader>
	void GetExports(const THeader* pNTHeader);

	template<class THeader>
	PIMAGE_COR20_HEADER GetCor20Header(const THeader* pNTHeader) const;

	template <class THeader> 
	LPVOID GetPtrFromRVA(DWORD rva, const THeader* pNTHeader) const;

	template <class THeader>
	PIMAGE_SECTION_HEADER GetEnclosingSectionHeader(DWORD rva, const THeader* pNTHeader) const;

	PBYTE                 m_base;
	CAtlFile              m_file;
	CAtlFileMapping<BYTE> m_map;

	bool m_Managed;
	bool m_Win32;
	bool m_COM;
	export_vector m_Exports;
};
